import numpy as np
import matplotlib.pyplot as plt

#X = np.arange(0,5,0.1)
#Z = [2+3 * x +4*x **2 for x in X]
#Y = np.array([np.random.normal(z,3) for z in Z])

Q = np.array([0.0,199.6,322.4,421.3,522.4,573.3,622.8,672.0,716.7,733.9,791.6,890.9,980.3])
H = np.array([153.6,155.0,154.95,154.91,152.82,150.38,148.13,145.34,143.25,142.03,138.99,131.68,123.97])



def gen_coefficient_matrix(X,Y):
    N = len(X)
    m =3
    A = []

    for i in range(m):
        a = []
        for j in range(m):
            a.append(sum(X**(i+j)))
        A.append(a)
    return A

def gen_right_vector(X,Y):
    N =len(X)
    m = 3
    b = []
    for i in range(m):
        b.append(sum(X**i * Y))
    return b
A = gen_coefficient_matrix(Q,H)
b = gen_right_vector(Q,H)

a0,a1,a2 = np.linalg.solve(A,b)

_X = np.arange(0,1000)
_Y = np.array([a0 +a1*x + a2*x**2 for x in _X])

plt.plot(Q,H,'ro', _X, _Y,'b',linewidth=2)
plt.ylim(ymax=200,ymin=50)
plt.title("y = {} + {}x + {}$x^2$".format(a0,a1,a2))
plt.annotate('H-Q', xy=(716.7,143.25), xytext=(200, 60), arrowprops=dict(facecolor='black', shrink=0.05),)
plt.show()